import numpy as np
import cv2
import time
import math

# window_size is adjusted according to image_size
# ================================================
def getDisp(imgL, imgR):
    """ Depth map calculation. Works with SGBM and WLS. Need rectified images, returns depth map ( left to right disparity ) """
    # SGBM Parameters -----------------
    window_size = 5  # wsize default 3; 5; 7 for SGBM reduced size image; 15 for SGBM full size image (1300px and above); 5 Works nicely

    left_matcher = cv2.StereoSGBM_create(
        minDisparity=-1,
        numDisparities=3*16,  # max_disp has to be dividable by 16 f. E. HH 192, 256
        blockSize=window_size,
        P1=8 * 3 * window_size*window_size,
        # wsize default 3; 5; 7 for SGBM reduced size image; 15 for SGBM full size image (1300px and above); 5 Works nicely
        P2=32 * 3 * window_size*window_size,
        disp12MaxDiff=1,
        uniquenessRatio=6, 
        speckleWindowSize=50, # 连通域大小
        speckleRange=1, # 连通域平滑阈值
        preFilterCap=15,
        mode=cv2.STEREO_SGBM_MODE_SGBM_3WAY
    )
    right_matcher = cv2.ximgproc.createRightMatcher(left_matcher)

    displ = left_matcher.compute(imgL, imgR)  # .astype(np.float32)/16
    dispr = right_matcher.compute(imgR, imgL)  # .astype(np.float32)/16
    displ = np.int16(displ)
    dispr = np.int16(dispr)

    lmbda = 80000
    sigma = 1.3
    visual_multiplier = 6

    wls_filterL = cv2.ximgproc.createDisparityWLSFilter(matcher_left=left_matcher)
    wls_filterL.setLambda(lmbda)
    wls_filterL.setSigmaColor(sigma)

    filteredImgL = wls_filterL.filter(displ, imgL, None, dispr)  # important to put "imgL" here!!!
    filteredImgL = filteredImgL/16
    filteredImgL8 = (filteredImgL).astype(np.uint8)
    return filteredImgL8



# ================================================
videoCapture = cv2.VideoCapture('newteststable.avi')  
success, frame = videoCapture.read()
width = int(0.5*frame.shape[1])
height = frame.shape[0]
# ================================================
# click to get depth
clicked = 0
clicked_x1 = 0
clicked_y1 = 0
clicked_x2 = 0
clicked_y2 = 0

focal = 3360
base_line = 200
u0 = 640
v0 = 360
intrinsic = [[focal, 0, u0],
             [0, focal, v0],
             [0,     0,  1]]
#print("intrinsic\n", intrinsic)
invIntrinsic = np.linalg.inv(intrinsic)
#print("invIntrinsic\n", invIntrinsic)


# mouse callback function
left = None
def get_pix_value(event,x,y,flags,param):
    if event == cv2.EVENT_LBUTTONDBLCLK:
        global clicked, clicked_x1, clicked_y1, clicked_x2, clicked_y2
        clicked = clicked + 1
        if clicked == 1:
            clicked_x1 = x
            clicked_y1 = y
        if clicked == 2:
            clicked_x2 = x
            clicked_y2 = y

        
# Create a window and bind the function to window
cv2.namedWindow('left')
cv2.setMouseCallback('left',get_pix_value)

# ================================================
cal_time = 0
# ================================================
def main():
    while True:
        start = time.time()
        success, frame = videoCapture.read()
        if not success:break
        left = frame[0:height, 0:width]
        right = frame[0:height, width:2*width]
        disp = getDisp(left,right)
        global clicked, clicked_x1, clicked_y1, clicked_x2, clicked_y2
        if clicked == 1:
            cv2.circle(left, (clicked_x1, clicked_y1), 5, (0, 255, 0), -1)
            zc1 = 0.001*(focal*base_line/disp[clicked_y1, clicked_x1])
            xc1 = zc1*(invIntrinsic[0,0]*clicked_x1 + 0 + invIntrinsic[0,2])
            yc1 = zc1*(0 + invIntrinsic[1,1]*clicked_y1 + invIntrinsic[1,2])
            #print("xc1,yc1,zc1 :  ", xc1,yc1,zc1)
        elif clicked == 2:
            cv2.circle(left, (clicked_x1, clicked_y1), 5, (255, 0, 0), -1)
            cv2.circle(left, (clicked_x2, clicked_y2), 5, (255, 0, 0), -1)
            cv2.line(left,(clicked_x1, clicked_y1),(clicked_x2, clicked_y2),(0,255,0),3)
            zc2 = 0.001*(focal*base_line/disp[clicked_y2, clicked_x2])
            xc2 = zc2*(invIntrinsic[0,0]*clicked_x2 + 0 + invIntrinsic[0,2])
            yc2 = zc2*(0 + invIntrinsic[1,1]*clicked_y2 + invIntrinsic[1,2])
            #print("xc2,yc2,zc2 :  ", xc2,yc2,zc2)
            length3d = math.sqrt((xc2-xc1)*(xc2-xc1) + (yc2-yc1)*(yc2-yc1) + (zc2-zc1)*(zc2-zc1))
            cv2.putText(left, str(length3d), (clicked_x2,clicked_y2), 0, 0.75, (0, 0, 255), 2)
        elif clicked == 3:
            clicked = 0
        
        disp_color=cv2.applyColorMap(cv2.convertScaleAbs(disp,alpha=2),cv2.COLORMAP_JET)
        cv2.imshow("left", left)
        #cv2.imshow("disp", disp)
        cv2.imshow("disp_color", disp_color)
        if cv2.waitKey(1)==27:break
        end = time.time()
        cal_time = end - start
        #print("FPS: ", 1/all_time)

main()













